import argparse
import backtrader as bt
from datetime import datetime as dt
from time import time
import multiprocessing
import os

from os.path import join, exists, dirname
if __name__ == '__main__':
    import sys
    sys.path.append(dirname(dirname(__file__)))
from strategies.netEasyLongMen import StrategyClass
from utils.localdata import onebyone, download

def main(num, ds, StrategyClass=StrategyClass):
    # for daily_data, code, name in [d for d in onebyone(klen=300) if d[1] in ['sh.600378', 'sh.603333']]:
    for kdata, code, name in ds:
        cerebro = bt.Cerebro()
        # 设置启动资金
        cerebro.broker.setcash(10000000.0)
        # 设置佣金
        cerebro.broker.setcommission(commission=0.001)
        cerebro.addstrategy(StrategyClass, code=f'[{num}] {code}', name=name, log=False)
        data = bt.feeds.PandasData(dataname=kdata)
        cerebro.adddata(data)
        cerebro.addsizer(bt.sizers.FixedSize, stake=100)  # 每笔交易使用固定交易量，stake代表交易量
        cerebro.run(plot=False)

def parse_args():
    from multiprocessing import cpu_count

    parser = argparse.ArgumentParser(description='鱼跃龙门形态搜索-多进程')

    parser.add_argument('-n', default=cpu_count()-1, required=False, type=int, help='进程数')

    parser.add_argument('-d', dest='download', action='store_const',
                        const=True, default=False, help='是否下载数据')

    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()

    if args.download:
        download()
        download('month')

    ds = [x for x in onebyone(frequency='month', klen=0, limitgt=61, types=['sh.60', 'sh.68', 'sz.00', 'sz.30'])]

    size = len(ds) // (args.n - 1)

    for i in range(args.n):
        multiprocessing.Process(target=main,kwargs={"num": i, "ds": ds[size * i:size * (i + 1)]}).start()

    print(dt.now().isoformat(), "主进程ID:", os.getpid(), "标的数量: ", len(ds), "每个进程搜索标的数量：", size)
